import torch
from torch import nn
from torch.nn import Conv2d,MaxPool2d,Flatten,Linear,Sequential,ReLU,Upsample,ConvTranspose2d
import torch.nn.functional as F
import numpy as np



class U_Net(nn.Module):
    def __init__(self):
        super(U_Net,self).__init__()

        self.conv11 = Conv2d(3,32,3,padding=1)
        self.relu = ReLU()
        self.conv12 = Conv2d(32,32,3,padding=1)
        self.maxpool = MaxPool2d(2)


        self.conv21 = Conv2d(32,64,3,padding=1)
        self.conv22 = Conv2d(64, 64, 3, padding=1)


        self.conv31 = Conv2d(64, 128, 3, padding=1)
        self.conv32 = Conv2d(128, 128, 3, padding=1)


        self.conv41 = Conv2d(128, 256, 3, padding=1)
        self.conv42 = Conv2d(256,256, 3, padding=1)

        self.conv51 = Conv2d(256, 512, 3, padding=1)
        self.conv52 = Conv2d(512, 512, 3, padding=1)


        self.upconv6 = ConvTranspose2d(512,256,kernel_size=2,stride=2,padding=0)       #stride=2 输出变成两倍    padding=2是自己算的
        #self.concat6 = np.concatenate([self.conv42,self.upconv6])
        #self.cat = torch.cat((),2)
        #self.concat6 = torch.cat((self.conv42,self.upconv6),2)
        self.conv61 = Conv2d(512, 256, 3, padding=1)
        self.conv62 = Conv2d(256, 256, 3, padding=1)

        self.upconv7 = ConvTranspose2d(256,128,kernel_size=2,stride=2,padding=0)       #stride=2 输出变成两倍    padding=2是自己算的
        #self.concat7 = np.concatenate([self.conv32,self.upconv7])
        #self.concat7 = torch.cat((self.conv32, self.upconv7), 2)
        self.conv71 = Conv2d(256, 128, 3, padding=1)
        self.conv72 = Conv2d(128, 128, 3, padding=1)

        self.upconv8 = ConvTranspose2d(128,64,kernel_size=2,stride=2,padding=0)       #stride=2 输出变成两倍    padding=2是自己算的
        #self.concat8 = np.concatenate([self.conv22,self.upconv8])
        #self.concat8 = torch.cat((self.conv22, self.upconv8), 2)
        self.conv81 = Conv2d(128, 64, 3, padding=1)
        self.conv82 = Conv2d(64, 64, 3, padding=1)

        self.upconv9 = ConvTranspose2d(64,32,kernel_size=2,stride=2,padding=0)       #stride=2 输出变成两倍    padding=2是自己算的
        #self.concat9 = np.concatenate([self.conv12,self.upconv9])
        #self.concat9 = torch.cat((self.conv12, self.upconv9), 2)
        self.conv91 = Conv2d(64, 32, 3, padding=1)
        self.conv92 = Conv2d(32, 32, 3, padding=1)

        self.conv93 = Conv2d(32,1,1,padding=0)


    def forward(self, x):
        x = self.conv11(x)
        x = self.relu(x)
        x = self.conv12(x)
        x1 = self.relu(x)
        x = self.maxpool(x1)

        x = self.conv21(x)
        x = self.relu(x)
        x = self.conv22(x)
        x2 = self.relu(x)
        x = self.maxpool(x2)

        x = self.conv31(x)
        x = self.relu(x)
        x = self.conv32(x)
        x3 = self.relu(x)
        x = self.maxpool(x3)

        x = self.conv41(x)
        x = self.relu(x)
        x = self.conv42(x)
        x4 = self.relu(x)
        x = self.maxpool(x4)

        x = self.conv51(x)
        x_train1 = self.relu(x)
        x = self.conv52(x_train1)
        x5 = self.relu(x)

        x = self.upconv6(x5)
        #x = self.concat6(x4,x)
        x = torch.cat((x4,x),1)
        x = self.conv61(x)
        x = self.relu(x)
        x = self.conv62(x)
        x = self.relu(x)

        x = self.upconv7(x)
        #x = self.concat7(x3,x)
        x = torch.cat((x3,x),1)
        x = self.conv71(x)
        x = self.relu(x)
        x = self.conv72(x)
        x = self.relu(x)

        x = self.upconv8(x)
        #x = self.concat8(x2,x)
        x = torch.cat((x2,x),1)
        x = self.conv81(x)
        x = self.relu(x)
        x = self.conv82(x)
        x = self.relu(x)

        x = self.upconv9(x)
        #x = self.concat9(x1,x)
        x = torch.cat((x1,x),1)
        x = self.conv91(x)
        x = self.relu(x)
        x = self.conv92(x)
        x = self.relu(x)

        x = self.conv93(x)
        return x 


